import java.text.DecimalFormat;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.text.*;
import java.util.*;
import java.io.*;

public class ThirdOrderHMM extends HMM {

  double[][][][] loga;
     
  private static double[][] toFirstOrder ( double[][][][] amat ) {
	double[][] aux = new double[amat.length][amat.length];
	for ( int i = 0 ; i < amat.length; i++ ) for ( int j = 0 ; j < amat.length ; j++ ) aux[i][j] = 0.0;
	for ( int i = 0 ; i < amat.length; i++ ) for ( int jj = 0 ; jj < amat.length ; jj++ ) for ( int j = 0 ; j < amat.length ; j++ ) for ( int k = 0 ; k < amat.length ; k++ ) aux[j][k] += amat[i][jj][j][k];
	return aux;
  }

  public ThirdOrderHMM(Object[] state, double[][][][] amat, Object[] esym, double[][] emat) {
    super(state,toFirstOrder(amat),esym,emat);
    if (state.length != amat.length) throw new IllegalArgumentException("HMM: state and amat disagree");
    if (amat.length != emat.length) throw new IllegalArgumentException("HMM: amat and emat disagree");
    for (int i=0; i<amat.length; i++) {
      if (state.length != amat[i].length) throw new IllegalArgumentException("HMM: amat non-square");
      if (esym.length != emat[i].length) throw new IllegalArgumentException("HMM: esym and emat disagree");
    }
    nstate = state.length + 1;
    this.state = new Object[nstate];
    loga = new double[nstate][nstate][nstate][nstate];
    for (int i=0; i<nstate; i++) for (int jj=0; jj<nstate; jj++) for (int j=0; j<nstate; j++) for (int k=0; k<nstate; k++) loga[i][jj][j][k] = 0.0;
    this.state[0] = (Object)(new String("---INITIAL-STATE---"));
    loga[0][0][0][0] = Double.NEGATIVE_INFINITY;
    double fromstart = Math.log(1.0 / state.length);
    for (int i=1; i<nstate; i++) loga[0][0][0][i] = fromstart;
    for (int i=1; i<nstate; i++) {
      this.state[i] = state[i-1];
      loga[i][0][0][0] = Double.NEGATIVE_INFINITY;
      loga[0][i][0][0] = Double.NEGATIVE_INFINITY;
      loga[0][0][i][0] = Double.NEGATIVE_INFINITY;
      for (int jj=1; jj<nstate; jj++) {
       loga[i][jj][0][0] = Double.NEGATIVE_INFINITY;
       loga[i][0][jj][0] = Double.NEGATIVE_INFINITY;
       loga[i][0][0][jj] = Double.NEGATIVE_INFINITY;
       loga[0][i][jj][0] = Double.NEGATIVE_INFINITY;
       loga[0][i][0][jj] = Double.NEGATIVE_INFINITY;
       for (int j=1; j<nstate; j++) {
         loga[i][jj][j][0] = Double.NEGATIVE_INFINITY;
         loga[i][jj][0][j] = Double.NEGATIVE_INFINITY;
         loga[i][0][jj][j] = Double.NEGATIVE_INFINITY;
         for (int k=1; k<nstate; k++) {
		loga[i][jj][j][k] = Math.log(amat[i-1][jj-1][j-1][k-1]);
	        loga[0][jj][j][k] += amat[i-1][jj-1][j-1][k-1];
	        loga[0][0][j][k] += amat[i-1][jj-1][j-1][k-1];
	 }
       }
      }
    }
    for (int j=1; j<nstate; j++) for (int k=1; k<nstate; k++) loga[0][0][j][k] = Math.log(loga[0][0][j][k]);
    for (int jj=1; jj<nstate; jj++) for (int j=1; j<nstate; j++) for (int k=1; k<nstate; k++) loga[0][jj][j][k] = Math.log(loga[0][jj][j][k]);
    this.esym = esym;
    nesym = esym.length;
    loge = new double[emat.length + 1][esym.length];
    for (int b=0; b < nesym; b++) {
      loge[0][b] = Double.NEGATIVE_INFINITY;
      for (int k=0; k<emat.length; k++) loge[k+1][b] = Math.log(emat[k][b]);
    }    
  }

  @Override
  public void printa(PrintStream out) {
    out.println("Transition probabilities:");
    for (int i=1; i<nstate; i++) {
      for (int jj=1; jj<nstate; jj++) for (int j=1; j<nstate; j++) for (int k=1; k<nstate; k++) out.print(fmtlog(loga[jj][j][k][i]));
      out.println();
    }
  }

  public static ThirdOrderHMM baumwelch(Object[][] xs, Object[] state, Object[] esym, final double threshold ) {
    int nstate = state.length;
    int nseqs  = xs.length;
    int nesym  = esym.length;
    ThirdOrderForward[] fwds = new ThirdOrderForward[nseqs];
    ThirdOrderBackward[] bwds = new ThirdOrderBackward[nseqs];
    double[] logP = new double[nseqs];
    double[][][][] amat = new double[nstate][nstate][nstate][];
    double[][] emat = new double[nstate][];
    int[] esyminv = new int[esym.length];
    for (int i=0; i<esyminv.length; i++) esyminv[i] = -1;
    for (int b=0; b<nesym; b++) esyminv[b] = b;
    for (int j=0; j<nstate; j++) for (int k=0; k<nstate; k++) for (int l=0; l<nstate; l++) amat[j][k][l] = uniformdiscrete(nstate);
    for (int k=0; k<nstate; k++) emat[k] = randomdiscrete(nesym);
    ThirdOrderHMM hmm = new ThirdOrderHMM(state, amat, esym, emat);  
    double oldloglikelihood;
    double loglikelihood = fwdbwd(hmm, xs, fwds, bwds, logP);
    int numIterations = 0;
    do {
      if ( maxIterationsBaumWelsh > -1 && ( numIterations++ > maxIterationsBaumWelsh ) ) break;
      oldloglikelihood = loglikelihood;
      double[][][][] A = new double[nstate][nstate][nstate][nstate];
      double[][] E = new double[nstate][nesym];
      for (int s=0; s < nseqs; s++) {
        Object[] x = xs[s];
        ThirdOrderForward fwd  = (ThirdOrderForward)(fwds[s]);
        ThirdOrderBackward bwd = (ThirdOrderBackward)(bwds[s]);
        int L = x.length;
        double P = logP[s];
        for (int i=0; i < L; i++) {
          for (int k=0; k<nstate; k++) E[k][esyminv[hmm.getSym(x[i])]] += exp(fwd.f[i+1][k+1] + bwd.b[i+1][k+1] - P);
        }
        for (int i=0; i<L-1; i++) 
          for (int jj=0; jj<nstate; jj++) for (int j=0; j<nstate; j++) 
	    for (int k=0; k<nstate; k++) 
	      for (int ell=0; ell<nstate; ell++) 
	              A[jj][j][k][ell] += exp(fwd.f[i+1][k+1] 
	                               + hmm.loga[jj+1][j+1][k+1][ell+1] 
	                               + hmm.loge[ell+1][hmm.getSym(x[i+1])] 
	                               + bwd.b[i+2][ell+1] 
	                               - P);
      }
      for (int jj=0; jj<nstate; jj++) for (int j=0; j<nstate; j++) for (int k=0; k<nstate; k++) {
        double Aksum = 0;
        for (int ell=0; ell<nstate; ell++) Aksum += A[jj][j][k][ell];
        for (int ell=0; ell<nstate; ell++) amat[jj][j][k][ell] = ( A[jj][j][k][ell] + 1.0 ) / ( Aksum + 0.0 + (nstate * nstate * nstate) );
      }
      for (int k=0; k<nstate; k++) {
        double Eksum = 0;
        for (int b=0; b<nesym; b++) Eksum += E[k][b];
        for (int b=0; b<nesym; b++) emat[k][b] = ( E[k][b] + 1.0 ) / ( Eksum + 0.0 + nesym);
      }
      hmm = new ThirdOrderHMM(state, amat, esym, emat);
      loglikelihood = fwdbwd(hmm, xs, fwds, bwds, logP);
      if ( debugBaumWelsh ) System.err.println("Baum-Welch Training - log likelihood = " + loglikelihood + " ( diff from previous is " + Math.abs(oldloglikelihood - loglikelihood) + ")");
    } while (Math.abs(oldloglikelihood - loglikelihood) > threshold);   
    return hmm;
  }

  protected static double fwdbwd(HMM hmm, Object[][] xs, Forward[] fwds, Backward[] bwds, double[] logP) {
    double loglikelihood = 0;
    for (int s=0; s<xs.length; s++) {
      fwds[s] = new ThirdOrderForward((ThirdOrderHMM)hmm, xs[s]);
      bwds[s] = new ThirdOrderBackward((ThirdOrderHMM)hmm, xs[s]);
      logP[s] = ((ThirdOrderForward)(fwds[s])).logprob();
      loglikelihood += logP[s];
    }
    return loglikelihood;
  }

  protected static double[] uniformdiscrete(int n) {
    double[] ps = new double[n];
    for (int i=0; i<n; i++) ps[i] = 1.0/(n * n * n);
    return ps;
  }

}
